import torch
import torch.nn as nn
from torch.nn import functional as F
class PPM(nn.Module):
    """
    Reference:
        Zhao, Hengshuang, et al. *"Pyramid scene parsing network."*
    """
    def __init__(self, in_channels, norm_layer=nn.BatchNorm2d):
        super(PPM, self).__init__()
        self.out_channels = in_channels*2
        self.pool1 = nn.AdaptiveAvgPool2d(1)
        self.pool2 = nn.AdaptiveAvgPool2d(2)
        self.pool3 = nn.AdaptiveAvgPool2d(3)
        self.pool4 = nn.AdaptiveAvgPool2d(6)
        out_channels = int(in_channels/4)
        self.conv1 = nn.Sequential(nn.Conv2d(in_channels, out_channels, 1, bias=False),
                                norm_layer(out_channels),
                                nn.ReLU(True))
        self.conv2 = nn.Sequential(nn.Conv2d(in_channels, out_channels, 1, bias=False),
                                norm_layer(out_channels),
                                nn.ReLU(True))
        self.conv3 = nn.Sequential(nn.Conv2d(in_channels, out_channels, 1, bias=False),
                                norm_layer(out_channels),
                                nn.ReLU(True))
        self.conv4 = nn.Sequential(nn.Conv2d(in_channels, out_channels, 1, bias=False),
                                norm_layer(out_channels),
                                nn.ReLU(True))
        self.conv5 = nn.Sequential(nn.Conv2d(in_channels*2, in_channels, 1, bias=False),
                                   norm_layer(in_channels),
                                   nn.ReLU(True))
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)

    def channels(self):
        return self.out_channels

    def forward(self, x):
        _, _, h, w = x.size()
        feat1 = F.interpolate(self.conv1(self.pool1(x)), (h, w))
        feat2 = F.interpolate(self.conv2(self.pool2(x)), (h, w))
        feat3 = F.interpolate(self.conv3(self.pool3(x)), (h, w))
        feat4 = F.interpolate(self.conv4(self.pool4(x)), (h, w))
        feat = torch.cat((x, feat1, feat2, feat3, feat4), 1)
        out = self.conv5(feat)
        out = self.pool(out)
        return out
